# -*- coding: utf-8 -*-

"""
统计词频
多线程
遍历文件夹下csv
"""

import os
import pandas as pd
import jieba
from concurrent.futures import ThreadPoolExecutor

from scripts.cut_word.thulac_cut_word import THULACSegmentor

# 加载自定义词
user_dict_path = "/home/ubuntu/code/git/subject-word-extraction/data/user_dict/"
thulac_segmentor = THULACSegmentor(user_dict_path=user_dict_path+'user_dict.txt')


def cut_content(content: str):
    """
    切割文本，剔除停用词。
    
    :param str content: 被处理的文本。
    :return: 切好的列表，其中每个元素是一个词语。
    :rtype: list[str]
    """
    cut = lambda x: [segment for segment in thulac_segmentor.segment(x)]
    segments = cut(content)
    
    word_list = []
    for segment in segments:
        if segment[1] in ['n', 'nt', 'nz', 'nl', 'v', 'vd', 'vn', 'vf', 'vx', 
                          'vi', 'vl', 'a', 'r', 'd', 'h', 'k'] and len(segment[0]) > 1:
            word_list.append(segment[0])
    
    return word_list

def calculate_word_frequency(word_list: list):
    """
    计算给定列表中的词语出现次数及其词频。
    
    :param list[str] word_list: 给定的词语列表。
    :return: 每个词语的出现次数及频率。
    :rtype: tuple[dict, dict]
    """
    n = len(word_list)
    word_count = dict()
    for word in word_list:
        if word in word_count.keys():
            word_count[word] += 1
        else:
            word_count[word] = 1
    
    word_frequency = dict()
    for word in word_count:
        word_frequency[word] = word_count[word]/n
    
    return word_count, word_frequency

def task(context, key_words, year, corp_name):
    """线程任务"""
    ls1 = []
    ls2 = []
    word_list = cut_content("".join(context))

    # 统计词频
    word_count, word_frequency = calculate_word_frequency(word_list)
    frequencies = dict()
    word_frequencies = dict()
    for key_word in key_words:
        frequencies[key_word] = []
        word_frequencies[key_word] = []

    for key_word in key_words:
        if key_word in word_count.keys():
            frequencies[key_word].append(word_count[key_word])
        else:
            frequencies[key_word].append(0)

        if key_word in word_frequency.keys():
            word_frequencies[key_word].append(word_frequency[key_word])
        else:
            word_frequencies[key_word].append(0)


    ls1.append(year)
    ls2.append(corp_name)
    tmp_df1 = pd.DataFrame({"year":ls1, "corp_name": ls2})
    tmp_df2 = pd.DataFrame(frequencies)
    tmp_df3 = pd.DataFrame(word_frequencies)
    df1 = pd.concat([tmp_df1, tmp_df2], axis=1)
    df2 = pd.concat([tmp_df1, tmp_df3], axis=1)
    print("\n{}".format(corp_name))

    return df1, df2


if __name__ == "__main__":

    # 关键词
    key_words = ['人工智能', '智能制造', '智慧制造', '主动制造','智能化转型','智能化','商业智能','图像理解',
                 '智能数据分析', '智能机器人', '制造执行系统', '智造', '机器学习', '深度学习', '一体化', '无人化',
                 '互联网技术', '工业互联网']

    folder_path = "/home/ubuntu/code/git/subject-word-extraction/data/output/"
    for root, _, files in os.walk(folder_path):
        for file in files:
            if file.endswith(".csv.gz"):
                file_path = os.path.join(folder_path, file)
                file_name = file.split(".")[0]

                # 读取文件
                df = pd.read_csv(file_path)

                # 创建线程池
                res_df = pd.DataFrame()
                futures = []
                with ThreadPoolExecutor(max_workers=8) as executor:
                    # 提交任务到线程池
                    for i in range(df.shape[0]):
                        print("执行 {}/{}".format(i+1, df.shape[0]), end="\r")
                        future = executor.submit(task, context=df['file_content'].values[i], key_words=key_words, year=df["year"].values[i], corp_name=df['corp_name'].values[i])
                        futures.append(future)

                # 获取线程返回结果
                res_df1 = pd.DataFrame()
                res_df2 = pd.DataFrame()
                for future in futures:
                    res_df1 = pd.concat([res_df1, future.result()[0]], axis=0)
                    res_df2 = pd.concat([res_df2, future.result()[1]], axis=0)
                res_df1.to_csv("/home/ubuntu/code/git/subject-word-extraction/data/out/{}thualac词结果.csv".format(file_name))
                res_df2.to_csv("/home/ubuntu/code/git/subject-word-extraction/data/out/{}thualac词频结果.csv".format(file_name))


